
#include "eye_tiling.h"
#include "register/op_def_registry.h"
#include "tiling_common.h"

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext *context) {

  EyeTilingData tiling;
  auto attrs = context->GetAttrs();
  auto r = attrs->GetInt(0);
  auto c = attrs->GetInt(1);
  auto cValue = *c;
  if (cValue == 0) {
    cValue = *r;
  }
  auto bPtr = attrs->GetListInt(2);
  int64_t b = 1;
  for (auto i = 0; i < bPtr->GetSize(); ++i) {
    b *= bPtr->GetData()[i];
  }
  tiling.set_b(b);
  tiling.set_r(*r);
  tiling.set_c(cValue);

  std::cout << "b = " << tiling.get_b() << std::endl;
  std::cout << "r = " << tiling.get_r() << std::endl;
  std::cout << "c = " << tiling.get_c() << std::endl;

  tiling_zero_workspace(context);
  tiling_save_to_buffer(tiling, context);

  return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext *context) {
  const gert::Shape *x1_shape = context->GetInputShape(0);
  gert::Shape *y_shape = context->GetOutputShape(0);
  *y_shape = *x1_shape;
  return GRAPH_SUCCESS;
}
} // namespace ge

namespace ops {
class Eye : public OpDef {
public:
  explicit Eye(const char *name) : OpDef(name) {
    this->Input("y")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
    this->Output("y")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
    this->Attr("num_rows").Int();
    this->Attr("num_columns").AttrType(OPTIONAL).Int(0);
    this->Attr("batch_shape").AttrType(OPTIONAL).ListInt({1});
    this->Attr("dtype").AttrType(OPTIONAL).Int(0);

    this->SetInferShape(ge::InferShape);

    this->AICore().SetTiling(optiling::TilingFunc);
    this->AICore().AddConfig("ascend310b");
  }
};

OP_ADD(Eye);
} // namespace ops
